import torch
from torch.utils.data import Dataset
from PIL import Image


class MyDataset(Dataset):
    def __init__(self, txt_path, transform=None, target_transform=None):
        txtPath = open(txt_path, 'r')
        imgs = []
        for line in txtPath:
            all_value = line.strip().split()
            imgs.append([all_value[0], int(all_value[1])])
            self.imgs = imgs
            self.transform = transform
            self.target_transform = target_transform

    def __getitem__(self, item):
        txtPath, label = self.imgs[item]
        img = Image.open(txtPath).convert('RGB')
        if self.transform is not None:
            img = self.transform(img)
        return img, label

    def __len__(self):
        return len(self.imgs)
